{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0b01c06a",
   "metadata": {},
   "source": [
    "# CLIQ-ai’21 Graph4NLP Tutorial Demo: Text Classification"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a12e3e7e",
   "metadata": {},
   "source": [
    "### Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "132dd151",
   "metadata": {},
   "source": [
    "In this tutorial demo, we will use the Graph4NLP library to build a GNN-based text classification model. The model consists of \n",
    "- graph construction module (e.g., dependency based static graph)\n",
    "- graph embedding module (e.g., Bi-Fuse GraphSAGE)\n",
    "- prediction module (e.g., graph pooling + MLP classifier)\n",
    "\n",
    "We will use the built-in module APIs to build the model, and evaluate it on the TREC dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9b3c0d5",
   "metadata": {},
   "source": [
    "### Environment setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b946a75",
   "metadata": {},
   "source": [
    "Please follow the instructions [here](https://github.com/graph4ai/graph4nlp_demo#environment-setup) to set up the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cd5f5762",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using backend: pytorch\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "import os\n",
    "import time\n",
    "import datetime\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.multiprocessing\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import yaml\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from graph4nlp.pytorch.datasets.trec import TrecDataset\n",
    "from graph4nlp.pytorch.modules.evaluation.accuracy import Accuracy\n",
    "from graph4nlp.pytorch.modules.graph_construction import (\n",
    "    ConstituencyBasedGraphConstruction,\n",
    "    DependencyBasedGraphConstruction,\n",
    "    IEBasedGraphConstruction,\n",
    "    NodeEmbeddingBasedGraphConstruction,\n",
    "    NodeEmbeddingBasedRefinedGraphConstruction,\n",
    ")\n",
    "from graph4nlp.pytorch.modules.graph_construction.embedding_construction import WordEmbedding\n",
    "from graph4nlp.pytorch.modules.graph_embedding import GAT, GGNN, GraphSAGE\n",
    "from graph4nlp.pytorch.modules.loss.general_loss import GeneralLoss\n",
    "from graph4nlp.pytorch.modules.prediction.classification.graph_classification import FeedForwardNN\n",
    "from graph4nlp.pytorch.modules.utils import constants as Constants\n",
    "from graph4nlp.pytorch.modules.utils.generic_utils import EarlyStopping, grid, to_cuda\n",
    "from graph4nlp.pytorch.modules.utils.logger import Logger"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "824bee58",
   "metadata": {},
   "source": [
    "### Build the text classifier"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99827b27",
   "metadata": {},
   "source": [
    "Let's first build the GNN-based text classifier which contains three major components including graph construction module, graph embedding module and graph prediction module. \n",
    "\n",
    "For graph construction module, the Graph4NLP library provides built-in APIs to support both static graph construction methods (e.g., `dependency graph`, `constituency graph`, `IE graph`) and dynamic graph construction methods (e.g., `node embedding based graph`, `node embedding based refined graph`). When calling the graph construction API, users should also specify the `embedding style` (e.g., word2vec, BiLSTM, BERT) to initialize the node/edge embeddings. Both single-token and multi-token node/edge graphs are supported.\n",
    "\n",
    "For graph embedding module, the Graph4NLP library provides builti-in APIs to support both `undirectional` and `bidirectinal` versions for common GNNs such as `GCN`, `GraphSAGE`, `GAT` and `GGNN`. \n",
    "\n",
    "For graph prediction module, the Graph4NLP library provides a high-level graph classification prediction module which consists of a graph pooling component (e.g., average pooling, max pooling) and a multilayer perceptron (MLP)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "669b55bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextClassifier(nn.Module):\n",
    "    def __init__(self, vocab, label_model, config):\n",
    "        super(TextClassifier, self).__init__()\n",
    "        self.config = config\n",
    "        self.vocab_model = vocab\n",
    "        self.label_model = label_model\n",
    "        self.graph_name = self.config[\"graph_construction_args\"][\"graph_construction_share\"][\n",
    "            \"graph_name\"\n",
    "        ]\n",
    "        embedding_style = {\n",
    "            \"single_token_item\": True if self.graph_name != \"ie\" else False,\n",
    "            \"emb_strategy\": config.get(\"emb_strategy\", \"w2v_bilstm\"),\n",
    "            \"num_rnn_layers\": 1,\n",
    "            \"bert_model_name\": config.get(\"bert_model_name\", \"bert-base-uncased\"),\n",
    "            \"bert_lower_case\": True,\n",
    "        }\n",
    "\n",
    "        assert not (\n",
    "            self.graph_name in (\"node_emb\", \"node_emb_refined\") and config[\"gnn\"] == \"gat\"\n",
    "        ), \"dynamic graph construction does not support GAT\"\n",
    "\n",
    "        use_edge_weight = False\n",
    "        if self.graph_name == \"dependency\":\n",
    "            self.graph_topology = DependencyBasedGraphConstruction(\n",
    "                embedding_style=embedding_style,\n",
    "                vocab=vocab.in_word_vocab,\n",
    "                hidden_size=config[\"num_hidden\"],\n",
    "                word_dropout=config[\"word_dropout\"],\n",
    "                rnn_dropout=config[\"rnn_dropout\"],\n",
    "                fix_word_emb=not config[\"no_fix_word_emb\"],\n",
    "                fix_bert_emb=not config.get(\"no_fix_bert_emb\", False),\n",
    "            )\n",
    "        elif self.graph_name == \"constituency\":\n",
    "            self.graph_topology = ConstituencyBasedGraphConstruction(\n",
    "                embedding_style=embedding_style,\n",
    "                vocab=vocab.in_word_vocab,\n",
    "                hidden_size=config[\"num_hidden\"],\n",
    "                word_dropout=config[\"word_dropout\"],\n",
    "                rnn_dropout=config[\"rnn_dropout\"],\n",
    "                fix_word_emb=not config[\"no_fix_word_emb\"],\n",
    "                fix_bert_emb=not config.get(\"no_fix_bert_emb\", False),\n",
    "            )\n",
    "        elif self.graph_name == \"ie\":\n",
    "            self.graph_topology = IEBasedGraphConstruction(\n",
    "                embedding_style=embedding_style,\n",
    "                vocab=vocab.in_word_vocab,\n",
    "                hidden_size=config[\"num_hidden\"],\n",
    "                word_dropout=config[\"word_dropout\"],\n",
    "                rnn_dropout=config[\"rnn_dropout\"],\n",
    "                fix_word_emb=not config[\"no_fix_word_emb\"],\n",
    "                fix_bert_emb=not config.get(\"no_fix_bert_emb\", False),\n",
    "            )\n",
    "        elif self.graph_name == \"node_emb\":\n",
    "            self.graph_topology = NodeEmbeddingBasedGraphConstruction(\n",
    "                vocab.in_word_vocab,\n",
    "                embedding_style,\n",
    "                sim_metric_type=config[\"gl_metric_type\"],\n",
    "                num_heads=config[\"gl_num_heads\"],\n",
    "                top_k_neigh=config[\"gl_top_k\"],\n",
    "                epsilon_neigh=config[\"gl_epsilon\"],\n",
    "                smoothness_ratio=config[\"gl_smoothness_ratio\"],\n",
    "                connectivity_ratio=config[\"gl_connectivity_ratio\"],\n",
    "                sparsity_ratio=config[\"gl_sparsity_ratio\"],\n",
    "                input_size=config[\"num_hidden\"],\n",
    "                hidden_size=config[\"gl_num_hidden\"],\n",
    "                fix_word_emb=not config[\"no_fix_word_emb\"],\n",
    "                fix_bert_emb=not config.get(\"no_fix_bert_emb\", False),\n",
    "                word_dropout=config[\"word_dropout\"],\n",
    "                rnn_dropout=config[\"rnn_dropout\"],\n",
    "            )\n",
    "            use_edge_weight = True\n",
    "        elif self.graph_name == \"node_emb_refined\":\n",
    "            self.graph_topology = NodeEmbeddingBasedRefinedGraphConstruction(\n",
    "                vocab.in_word_vocab,\n",
    "                embedding_style,\n",
    "                config[\"init_adj_alpha\"],\n",
    "                sim_metric_type=config[\"gl_metric_type\"],\n",
    "                num_heads=config[\"gl_num_heads\"],\n",
    "                top_k_neigh=config[\"gl_top_k\"],\n",
    "                epsilon_neigh=config[\"gl_epsilon\"],\n",
    "                smoothness_ratio=config[\"gl_smoothness_ratio\"],\n",
    "                connectivity_ratio=config[\"gl_connectivity_ratio\"],\n",
    "                sparsity_ratio=config[\"gl_sparsity_ratio\"],\n",
    "                input_size=config[\"num_hidden\"],\n",
    "                hidden_size=config[\"gl_num_hidden\"],\n",
    "                fix_word_emb=not config[\"no_fix_word_emb\"],\n",
    "                fix_bert_emb=not config.get(\"no_fix_bert_emb\", False),\n",
    "                word_dropout=config[\"word_dropout\"],\n",
    "                rnn_dropout=config[\"rnn_dropout\"],\n",
    "            )\n",
    "            use_edge_weight = True\n",
    "        else:\n",
    "            raise RuntimeError(\"Unknown graph_name: {}\".format(self.graph_name))\n",
    "\n",
    "        if \"w2v\" in self.graph_topology.embedding_layer.word_emb_layers:\n",
    "            self.word_emb = self.graph_topology.embedding_layer.word_emb_layers[\n",
    "                \"w2v\"\n",
    "            ].word_emb_layer\n",
    "        else:\n",
    "            self.word_emb = WordEmbedding(\n",
    "                self.vocab_model.in_word_vocab.embeddings.shape[0],\n",
    "                self.vocab_model.in_word_vocab.embeddings.shape[1],\n",
    "                pretrained_word_emb=self.vocab_model.in_word_vocab.embeddings,\n",
    "                fix_emb=not config[\"no_fix_word_emb\"],\n",
    "            ).word_emb_layer\n",
    "\n",
    "        if config[\"gnn\"] == \"gat\":\n",
    "            heads = [config[\"gat_num_heads\"]] * (config[\"gnn_num_layers\"] - 1) + [\n",
    "                config[\"gat_num_out_heads\"]\n",
    "            ]\n",
    "            self.gnn = GAT(\n",
    "                config[\"gnn_num_layers\"],\n",
    "                config[\"num_hidden\"],\n",
    "                config[\"num_hidden\"],\n",
    "                config[\"num_hidden\"],\n",
    "                heads,\n",
    "                direction_option=config[\"gnn_direction_option\"],\n",
    "                feat_drop=config[\"gnn_dropout\"],\n",
    "                attn_drop=config[\"gat_attn_dropout\"],\n",
    "                negative_slope=config[\"gat_negative_slope\"],\n",
    "                residual=config[\"gat_residual\"],\n",
    "                activation=F.elu,\n",
    "                allow_zero_in_degree=True,\n",
    "            )\n",
    "        elif config[\"gnn\"] == \"graphsage\":\n",
    "            self.gnn = GraphSAGE(\n",
    "                config[\"gnn_num_layers\"],\n",
    "                config[\"num_hidden\"],\n",
    "                config[\"num_hidden\"],\n",
    "                config[\"num_hidden\"],\n",
    "                config[\"graphsage_aggreagte_type\"],\n",
    "                direction_option=config[\"gnn_direction_option\"],\n",
    "                feat_drop=config[\"gnn_dropout\"],\n",
    "                bias=True,\n",
    "                norm=None,\n",
    "                activation=F.relu,\n",
    "                use_edge_weight=use_edge_weight,\n",
    "            )\n",
    "        elif config[\"gnn\"] == \"ggnn\":\n",
    "            self.gnn = GGNN(\n",
    "                config[\"gnn_num_layers\"],\n",
    "                config[\"num_hidden\"],\n",
    "                config[\"num_hidden\"],\n",
    "                config[\"num_hidden\"],\n",
    "                feat_drop=config[\"gnn_dropout\"],\n",
    "                direction_option=config[\"gnn_direction_option\"],\n",
    "                bias=True,\n",
    "                use_edge_weight=use_edge_weight,\n",
    "            )\n",
    "        else:\n",
    "            raise RuntimeError(\"Unknown gnn type: {}\".format(config[\"gnn\"]))\n",
    "\n",
    "        self.clf = FeedForwardNN(\n",
    "            2 * config[\"num_hidden\"]\n",
    "            if config[\"gnn_direction_option\"] == \"bi_sep\"\n",
    "            else config[\"num_hidden\"],\n",
    "            config[\"num_classes\"],\n",
    "            [config[\"num_hidden\"]],\n",
    "            graph_pool_type=config[\"graph_pooling\"],\n",
    "            dim=config[\"num_hidden\"],\n",
    "            use_linear_proj=config[\"max_pool_linear_proj\"],\n",
    "        )\n",
    "\n",
    "        self.loss = GeneralLoss(\"CrossEntropy\")\n",
    "\n",
    "    def forward(self, graph_list, tgt=None, require_loss=True):\n",
    "        # graph construction\n",
    "        batch_gd = self.graph_topology(graph_list)\n",
    "\n",
    "        # run GNN\n",
    "        self.gnn(batch_gd)\n",
    "\n",
    "        # run graph classifier\n",
    "        self.clf(batch_gd)\n",
    "        logits = batch_gd.graph_attributes[\"logits\"]\n",
    "\n",
    "        if require_loss:\n",
    "            loss = self.loss(logits, tgt)\n",
    "            return logits, loss\n",
    "        else:\n",
    "            return logits\n",
    "\n",
    "    def inference_forward(self, collate_data):\n",
    "        return self.forward(collate_data[\"graph_data\"], require_loss=False)\n",
    "\n",
    "    def post_process(self, logits, label_names):\n",
    "        logits_list = []\n",
    "\n",
    "        for idx in range(len(logits)):\n",
    "            logits_list.append(logits[idx].cpu().clone().numpy())\n",
    "\n",
    "        pred_tags = [label_names[pred.argmax()] for pred in logits_list]\n",
    "        return pred_tags\n",
    "\n",
    "    @classmethod\n",
    "    def load_checkpoint(cls, model_path):\n",
    "        \"\"\"The API to load the model.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        model_path : str\n",
    "            The saved model path.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Class\n",
    "        \"\"\"\n",
    "        return torch.load(model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f4a64e2",
   "metadata": {},
   "source": [
    "### Build the model handler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07b8ff8b",
   "metadata": {},
   "source": [
    "Next, let's build a model handler which will do a bunch of things including setting up dataloader, model, optimizer, evaluation metrics, train/val/test loops, and so on.\n",
    "\n",
    "When setting up the dataloader, users will need to call the dataset API which will preprocess the data, e.g., calling the graph construction module, building the vocabulary, tensorizing the data. Users will need to specify the graph construction type when calling the dataset API.\n",
    "\n",
    "Users can build their customized dataset APIs by inheriting our low-level dataset APIs. We provide low-level dataset APIs to support various scenarios (e.g., `Text2Label`, `Sequence2Labeling`, `Text2Text`, `Text2Tree`, `DoubleText2Text`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eda1f453",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelHandler:\n",
    "    def __init__(self, config):\n",
    "        super(ModelHandler, self).__init__()\n",
    "        self.config = config\n",
    "        self.logger = Logger(\n",
    "            self.config[\"out_dir\"],\n",
    "            config={k: v for k, v in self.config.items() if k != \"device\"},\n",
    "            overwrite=True,\n",
    "        )\n",
    "        self.logger.write(self.config[\"out_dir\"])\n",
    "        self._build_dataloader()\n",
    "        self._build_model()\n",
    "        self._build_optimizer()\n",
    "        self._build_evaluation()\n",
    "\n",
    "    def _build_dataloader(self):\n",
    "        self.graph_name = self.config[\"graph_construction_args\"][\"graph_construction_share\"][\n",
    "            \"graph_name\"\n",
    "        ]\n",
    "        topology_subdir = \"{}_graph\".format(self.graph_name)\n",
    "        if self.graph_name == \"node_emb_refined\":\n",
    "            topology_subdir += \"_{}\".format(\n",
    "                self.config[\"graph_construction_args\"][\"graph_construction_private\"][\n",
    "                    \"dynamic_init_graph_name\"\n",
    "                ]\n",
    "            )\n",
    "\n",
    "        dataset = TrecDataset(\n",
    "            root_dir=self.config[\"graph_construction_args\"][\"graph_construction_share\"][\"root_dir\"],\n",
    "            topology_subdir=topology_subdir,\n",
    "            graph_name=self.graph_name,\n",
    "            dynamic_init_graph_name=self.config[\"graph_construction_args\"][\n",
    "                \"graph_construction_private\"\n",
    "            ][\"dynamic_init_graph_name\"],\n",
    "            dynamic_init_topology_aux_args={\"dummy_param\": 0},\n",
    "            pretrained_word_emb_name=self.config[\"pretrained_word_emb_name\"],\n",
    "            merge_strategy=self.config[\"graph_construction_args\"][\"graph_construction_private\"][\n",
    "                \"merge_strategy\"\n",
    "            ],\n",
    "            edge_strategy=self.config[\"graph_construction_args\"][\"graph_construction_private\"][\n",
    "                \"edge_strategy\"\n",
    "            ],\n",
    "            min_word_vocab_freq=self.config.get(\"min_word_freq\", 1),\n",
    "            word_emb_size=self.config.get(\"word_emb_size\", 300),\n",
    "            seed=self.config[\"seed\"],\n",
    "            thread_number=self.config[\"graph_construction_args\"][\"graph_construction_share\"][\n",
    "                \"thread_number\"\n",
    "            ],\n",
    "            port=self.config[\"graph_construction_args\"][\"graph_construction_share\"][\"port\"],\n",
    "            timeout=self.config[\"graph_construction_args\"][\"graph_construction_share\"][\"timeout\"],\n",
    "            reused_label_model=None,\n",
    "        )\n",
    "\n",
    "        self.train_dataloader = DataLoader(\n",
    "            dataset.train,\n",
    "            batch_size=self.config[\"batch_size\"],\n",
    "            shuffle=True,\n",
    "            num_workers=self.config[\"num_workers\"],\n",
    "            collate_fn=dataset.collate_fn,\n",
    "        )\n",
    "        if not hasattr(dataset, \"val\"):\n",
    "            dataset.val = dataset.test\n",
    "        self.val_dataloader = DataLoader(\n",
    "            dataset.val,\n",
    "            batch_size=self.config[\"batch_size\"],\n",
    "            shuffle=False,\n",
    "            num_workers=self.config[\"num_workers\"],\n",
    "            collate_fn=dataset.collate_fn,\n",
    "        )\n",
    "        self.test_dataloader = DataLoader(\n",
    "            dataset.test,\n",
    "            batch_size=self.config[\"batch_size\"],\n",
    "            shuffle=False,\n",
    "            num_workers=self.config[\"num_workers\"],\n",
    "            collate_fn=dataset.collate_fn,\n",
    "        )\n",
    "        self.vocab_model = dataset.vocab_model\n",
    "        self.label_model = dataset.label_model\n",
    "        self.config[\"num_classes\"] = self.label_model.num_classes\n",
    "        self.num_train = len(dataset.train)\n",
    "        self.num_val = len(dataset.val)\n",
    "        self.num_test = len(dataset.test)\n",
    "        print(\n",
    "            \"Train size: {}, Val size: {}, Test size: {}\".format(\n",
    "                self.num_train, self.num_val, self.num_test\n",
    "            )\n",
    "        )\n",
    "        self.logger.write(\n",
    "            \"Train size: {}, Val size: {}, Test size: {}\".format(\n",
    "                self.num_train, self.num_val, self.num_test\n",
    "            )\n",
    "        )\n",
    "\n",
    "    def _build_model(self):\n",
    "        self.model = TextClassifier(self.vocab_model, self.label_model, self.config).to(\n",
    "            self.config[\"device\"]\n",
    "        )\n",
    "\n",
    "    def _build_optimizer(self):\n",
    "        parameters = [p for p in self.model.parameters() if p.requires_grad]\n",
    "        self.optimizer = optim.Adam(parameters, lr=self.config[\"lr\"])\n",
    "        self.stopper = EarlyStopping(\n",
    "            os.path.join(\n",
    "                self.config[\"out_dir\"],\n",
    "                self.config.get(\"model_ckpt_name\", Constants._SAVED_WEIGHTS_FILE),\n",
    "            ),\n",
    "            patience=self.config[\"patience\"],\n",
    "        )\n",
    "        self.scheduler = ReduceLROnPlateau(\n",
    "            self.optimizer,\n",
    "            mode=\"max\",\n",
    "            factor=self.config[\"lr_reduce_factor\"],\n",
    "            patience=self.config[\"lr_patience\"],\n",
    "            verbose=True,\n",
    "        )\n",
    "\n",
    "    def _build_evaluation(self):\n",
    "        self.metric = Accuracy([\"accuracy\"])\n",
    "\n",
    "    def train(self):\n",
    "        dur = []\n",
    "        for epoch in range(self.config[\"epochs\"]):\n",
    "            self.model.train()\n",
    "            train_loss = []\n",
    "            train_acc = []\n",
    "            t0 = time.time()\n",
    "            for data in self.train_dataloader:\n",
    "                tgt = to_cuda(data[\"tgt_tensor\"], self.config[\"device\"])\n",
    "                data[\"graph_data\"] = data[\"graph_data\"].to(self.config[\"device\"])\n",
    "                logits, loss = self.model(data[\"graph_data\"], tgt, require_loss=True)\n",
    "\n",
    "                # add graph regularization loss if available\n",
    "                if data[\"graph_data\"].graph_attributes.get(\"graph_reg\", None) is not None:\n",
    "                    loss = loss + data[\"graph_data\"].graph_attributes[\"graph_reg\"]\n",
    "\n",
    "                self.optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                self.optimizer.step()\n",
    "                train_loss.append(loss.item())\n",
    "\n",
    "                pred = torch.max(logits, dim=-1)[1].cpu()\n",
    "                train_acc.append(\n",
    "                    self.metric.calculate_scores(ground_truth=tgt.cpu(), predict=pred.cpu())[0]\n",
    "                )\n",
    "                dur.append(time.time() - t0)\n",
    "\n",
    "            val_acc = self.evaluate(self.val_dataloader)\n",
    "            self.scheduler.step(val_acc)\n",
    "            print(\n",
    "                \"Epoch: [{} / {}] | Time: {:.2f}s | Loss: {:.4f} |\"\n",
    "                \"Train Acc: {:.4f} | Val Acc: {:.4f}\".format(\n",
    "                    epoch + 1,\n",
    "                    self.config[\"epochs\"],\n",
    "                    np.mean(dur),\n",
    "                    np.mean(train_loss),\n",
    "                    np.mean(train_acc),\n",
    "                    val_acc,\n",
    "                )\n",
    "            )\n",
    "            self.logger.write(\n",
    "                \"Epoch: [{} / {}] | Time: {:.2f}s | Loss: {:.4f} |\"\n",
    "                \"Train Acc: {:.4f} | Val Acc: {:.4f}\".format(\n",
    "                    epoch + 1,\n",
    "                    self.config[\"epochs\"],\n",
    "                    np.mean(dur),\n",
    "                    np.mean(train_loss),\n",
    "                    np.mean(train_acc),\n",
    "                    val_acc,\n",
    "                )\n",
    "            )\n",
    "\n",
    "            if self.stopper.step(val_acc, self.model):\n",
    "                break\n",
    "\n",
    "        return self.stopper.best_score\n",
    "\n",
    "    def evaluate(self, dataloader):\n",
    "        self.model.eval()\n",
    "        with torch.no_grad():\n",
    "            pred_collect = []\n",
    "            gt_collect = []\n",
    "            for data in dataloader:\n",
    "                tgt = to_cuda(data[\"tgt_tensor\"], self.config[\"device\"])\n",
    "                data[\"graph_data\"] = data[\"graph_data\"].to(self.config[\"device\"])\n",
    "                logits = self.model(data[\"graph_data\"], require_loss=False)\n",
    "                pred_collect.append(logits)\n",
    "                gt_collect.append(tgt)\n",
    "\n",
    "            pred_collect = torch.max(torch.cat(pred_collect, 0), dim=-1)[1].cpu()\n",
    "            gt_collect = torch.cat(gt_collect, 0).cpu()\n",
    "            score = self.metric.calculate_scores(ground_truth=gt_collect, predict=pred_collect)[0]\n",
    "\n",
    "            return score\n",
    "\n",
    "    def test(self):\n",
    "        # restored best saved model\n",
    "        self.model = TextClassifier.load_checkpoint(self.stopper.save_model_path)\n",
    "\n",
    "        t0 = time.time()\n",
    "        acc = self.evaluate(self.test_dataloader)\n",
    "        dur = time.time() - t0\n",
    "        print(\n",
    "            \"Test examples: {} | Time: {:.2f}s |  Test Acc: {:.4f}\".format(self.num_test, dur, acc)\n",
    "        )\n",
    "        self.logger.write(\n",
    "            \"Test examples: {} | Time: {:.2f}s |  Test Acc: {:.4f}\".format(self.num_test, dur, acc)\n",
    "        )\n",
    "\n",
    "        return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ad3d820",
   "metadata": {},
   "source": [
    "### Set up the config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b2bff971",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_config(config):\n",
    "    print('**************** MODEL CONFIGURATION ****************')\n",
    "    for key in sorted(config.keys()):\n",
    "        val = config[key]\n",
    "        keystr = '{}'.format(key) + (' ' * (24 - len(key)))\n",
    "        print('{} -->   {}'.format(keystr, val))\n",
    "    print('**************** MODEL CONFIGURATION ****************')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "31b9931a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**************** MODEL CONFIGURATION ****************\n",
      "batch_size               -->   50\n",
      "dataset                  -->   trec\n",
      "emb_strategy             -->   w2v_bilstm\n",
      "epochs                   -->   500\n",
      "gat_attn_dropout         -->   None\n",
      "gat_negative_slope       -->   None\n",
      "gat_num_heads            -->   None\n",
      "gat_num_out_heads        -->   None\n",
      "gat_residual             -->   False\n",
      "gl_connectivity_ratio    -->   None\n",
      "gl_epsilon               -->   None\n",
      "gl_metric_type           -->   None\n",
      "gl_num_heads             -->   1\n",
      "gl_num_hidden            -->   300\n",
      "gl_smoothness_ratio      -->   None\n",
      "gl_sparsity_ratio        -->   None\n",
      "gl_top_k                 -->   None\n",
      "gnn                      -->   graphsage\n",
      "gnn_direction_option     -->   bi_fuse\n",
      "gnn_dropout              -->   0.4\n",
      "gnn_num_layers           -->   1\n",
      "gpu                      -->   -1\n",
      "graph_construction_args  -->   {'graph_construction_share': {'graph_name': 'dependency', 'root_dir': '../data/trec', 'thread_number': 10, 'port': 9000, 'timeout': 15000}, 'graph_construction_private': {'edge_strategy': 'homogeneous', 'merge_strategy': 'tailhead', 'sequential_link': True, 'as_node': False, 'dynamic_init_graph_name': None}}\n",
      "graph_pooling            -->   avg_pool\n",
      "graphsage_aggreagte_type -->   lstm\n",
      "init_adj_alpha           -->   None\n",
      "lr                       -->   0.001\n",
      "lr_patience              -->   2\n",
      "lr_reduce_factor         -->   0.5\n",
      "max_pool_linear_proj     -->   False\n",
      "no_cuda                  -->   False\n",
      "no_fix_word_emb          -->   False\n",
      "num_hidden               -->   300\n",
      "num_workers              -->   0\n",
      "out_dir                  -->   out/trec/graphsage_bi_fuse_dependency_ckpt\n",
      "patience                 -->   10\n",
      "pretrained_word_emb_name -->   840B\n",
      "rnn_dropout              -->   0.1\n",
      "seed                     -->   1234\n",
      "val_split_ratio          -->   0.2\n",
      "word_dropout             -->   0.4\n",
      "**************** MODEL CONFIGURATION ****************\n"
     ]
    }
   ],
   "source": [
    "# config setup\n",
    "config_file = '../config/trec/graphsage_bi_fuse_static_dependency_v2.yaml'\n",
    "config = yaml.load(open(config_file, 'r'), Loader=yaml.FullLoader)\n",
    "print_config(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ed47302",
   "metadata": {},
   "source": [
    "### Run the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a5087ca0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 545\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 545\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 546\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 545\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 545\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 545\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 545\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 546\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 545\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 545\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n",
      "Using backend: pytorch\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Connecting to stanfordcorenlp server...\n",
      "CoreNLP server connected.\n",
      "Port 9000, processing: 0 / 50\n",
      "Building vocabs...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      ".vector_cache/glove.840B.300d.zip: 2.18GB [12:51, 2.82MB/s]                                    \n",
      "100%|█████████████████████████████████████████████▉| 2196016/2196017 [05:07<00:00, 7134.20it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pretrained word embeddings hit ratio: 0.9572679106196729\n",
      "Using pretrained word embeddings\n",
      "[ Initialized word embeddings: (8682, 300) ]\n",
      "Saving vocab model to ../data/trec/processed/dependency_graph/vocab.pt\n",
      "Saving label mappings to ../data/trec/processed/dependency_graph/label.pt\n",
      "Loading pre-built label mappings stored in ../data/trec/processed/dependency_graph/label.pt\n",
      "Train size: 5452, Val size: 500, Test size: 500\n",
      "[ Fix word embeddings ]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hugo/opt/anaconda3/envs/g4l_nov2021b/lib/python3.7/site-packages/torch-1.10.0-py3.7-macosx-10.9-x86_64.egg/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling\n",
      "  warnings.warn('User provided device_type of \\'cuda\\', but CUDA is not available. Disabling')\n",
      "/Users/hugo/opt/anaconda3/envs/g4l_nov2021b/lib/python3.7/site-packages/graph4nlp-0.5.1-py3.7.egg/graph4nlp/pytorch/modules/evaluation/accuracy.py:157: UserWarning: zero division encountered\n",
      "  warnings.warn(\"zero division encountered\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [1 / 500] | Time: 13.19s | Loss: 1.1069 |Train Acc: 0.5576 | Val Acc: 0.7840\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch: [2 / 500] | Time: 13.13s | Loss: 0.6396 |Train Acc: 0.7647 | Val Acc: 0.7940\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch: [3 / 500] | Time: 13.10s | Loss: 0.4928 |Train Acc: 0.8278 | Val Acc: 0.8820\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch: [4 / 500] | Time: 13.13s | Loss: 0.4217 |Train Acc: 0.8569 | Val Acc: 0.8720\n",
      "EarlyStopping counter: 1 out of 10\n",
      "Epoch: [5 / 500] | Time: 13.15s | Loss: 0.3589 |Train Acc: 0.8700 | Val Acc: 0.9240\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch: [6 / 500] | Time: 13.15s | Loss: 0.3236 |Train Acc: 0.8873 | Val Acc: 0.9200\n",
      "EarlyStopping counter: 1 out of 10\n",
      "Epoch: [7 / 500] | Time: 13.15s | Loss: 0.2929 |Train Acc: 0.8973 | Val Acc: 0.9220\n",
      "EarlyStopping counter: 2 out of 10\n",
      "Epoch     8: reducing learning rate of group 0 to 5.0000e-04.\n",
      "Epoch: [8 / 500] | Time: 13.14s | Loss: 0.2344 |Train Acc: 0.9202 | Val Acc: 0.9220\n",
      "EarlyStopping counter: 3 out of 10\n",
      "Epoch: [9 / 500] | Time: 13.13s | Loss: 0.1853 |Train Acc: 0.9340 | Val Acc: 0.9380\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch: [10 / 500] | Time: 13.12s | Loss: 0.1692 |Train Acc: 0.9422 | Val Acc: 0.9280\n",
      "EarlyStopping counter: 1 out of 10\n",
      "Epoch: [11 / 500] | Time: 13.13s | Loss: 0.1341 |Train Acc: 0.9500 | Val Acc: 0.9260\n",
      "EarlyStopping counter: 2 out of 10\n",
      "Epoch    12: reducing learning rate of group 0 to 2.5000e-04.\n",
      "Epoch: [12 / 500] | Time: 13.12s | Loss: 0.1453 |Train Acc: 0.9471 | Val Acc: 0.9320\n",
      "EarlyStopping counter: 3 out of 10\n",
      "Epoch: [13 / 500] | Time: 13.12s | Loss: 0.1073 |Train Acc: 0.9622 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 4 out of 10\n",
      "Epoch: [14 / 500] | Time: 13.13s | Loss: 0.1084 |Train Acc: 0.9620 | Val Acc: 0.9280\n",
      "EarlyStopping counter: 5 out of 10\n",
      "Epoch    15: reducing learning rate of group 0 to 1.2500e-04.\n",
      "Epoch: [15 / 500] | Time: 13.13s | Loss: 0.1038 |Train Acc: 0.9627 | Val Acc: 0.9320\n",
      "EarlyStopping counter: 6 out of 10\n",
      "Epoch: [16 / 500] | Time: 13.14s | Loss: 0.0898 |Train Acc: 0.9669 | Val Acc: 0.9340\n",
      "EarlyStopping counter: 7 out of 10\n",
      "Epoch: [17 / 500] | Time: 13.14s | Loss: 0.0753 |Train Acc: 0.9720 | Val Acc: 0.9380\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch    18: reducing learning rate of group 0 to 6.2500e-05.\n",
      "Epoch: [18 / 500] | Time: 13.14s | Loss: 0.0843 |Train Acc: 0.9707 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 1 out of 10\n",
      "Epoch: [19 / 500] | Time: 13.15s | Loss: 0.0723 |Train Acc: 0.9749 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 2 out of 10\n",
      "Epoch: [20 / 500] | Time: 13.15s | Loss: 0.0725 |Train Acc: 0.9751 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 3 out of 10\n",
      "Epoch    21: reducing learning rate of group 0 to 3.1250e-05.\n",
      "Epoch: [21 / 500] | Time: 13.16s | Loss: 0.0699 |Train Acc: 0.9738 | Val Acc: 0.9380\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch: [22 / 500] | Time: 13.16s | Loss: 0.0729 |Train Acc: 0.9744 | Val Acc: 0.9380\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch: [23 / 500] | Time: 13.16s | Loss: 0.0663 |Train Acc: 0.9784 | Val Acc: 0.9400\n",
      "Saved model to out/trec/graphsage_bi_fuse_dependency_ckpt_1637304424.77838/params.saved\n",
      "Epoch: [24 / 500] | Time: 13.17s | Loss: 0.0673 |Train Acc: 0.9751 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 1 out of 10\n",
      "Epoch: [25 / 500] | Time: 13.17s | Loss: 0.0676 |Train Acc: 0.9744 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 2 out of 10\n",
      "Epoch    26: reducing learning rate of group 0 to 1.5625e-05.\n",
      "Epoch: [26 / 500] | Time: 13.17s | Loss: 0.0663 |Train Acc: 0.9758 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 3 out of 10\n",
      "Epoch: [27 / 500] | Time: 13.17s | Loss: 0.0586 |Train Acc: 0.9782 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 4 out of 10\n",
      "Epoch: [28 / 500] | Time: 13.17s | Loss: 0.0595 |Train Acc: 0.9762 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 5 out of 10\n",
      "Epoch    29: reducing learning rate of group 0 to 7.8125e-06.\n",
      "Epoch: [29 / 500] | Time: 13.17s | Loss: 0.0649 |Train Acc: 0.9785 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 6 out of 10\n",
      "Epoch: [30 / 500] | Time: 13.17s | Loss: 0.0594 |Train Acc: 0.9802 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 7 out of 10\n",
      "Epoch: [31 / 500] | Time: 13.17s | Loss: 0.0650 |Train Acc: 0.9755 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 8 out of 10\n",
      "Epoch    32: reducing learning rate of group 0 to 3.9063e-06.\n",
      "Epoch: [32 / 500] | Time: 13.17s | Loss: 0.0609 |Train Acc: 0.9780 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 9 out of 10\n",
      "Epoch: [33 / 500] | Time: 13.18s | Loss: 0.0675 |Train Acc: 0.9755 | Val Acc: 0.9360\n",
      "EarlyStopping counter: 10 out of 10\n",
      "Test examples: 500 | Time: 0.67s |  Test Acc: 0.9400\n",
      "Total runtime: 879.95s\n",
      "val acc: 0.94, test acc: 0.94\n"
     ]
    }
   ],
   "source": [
    "# run model\n",
    "import platform, multiprocessing\n",
    "\n",
    "if platform.system() == \"Darwin\": # MacOS\n",
    "    multiprocessing.set_start_method('spawn')\n",
    "\n",
    "\n",
    "np.random.seed(config['seed'])\n",
    "torch.manual_seed(config['seed'])\n",
    "\n",
    "if not config[\"no_cuda\"] and torch.cuda.is_available():\n",
    "    print(\"[ Using CUDA ]\")\n",
    "    config[\"device\"] = torch.device(\"cuda\" if config[\"gpu\"] < 0 else \"cuda:%d\" % config[\"gpu\"])\n",
    "    torch.cuda.manual_seed(config[\"seed\"])\n",
    "    torch.cuda.manual_seed_all(config[\"seed\"])\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    cudnn.benchmark = False\n",
    "else:\n",
    "    config[\"device\"] = torch.device(\"cpu\")\n",
    "\n",
    "\n",
    "ts = datetime.datetime.now().timestamp()\n",
    "config['out_dir'] += '_{}'.format(ts)\n",
    "print('\\n' + config['out_dir'])\n",
    "\n",
    "runner = ModelHandler(config)\n",
    "t0 = time.time()\n",
    "\n",
    "val_acc = runner.train()\n",
    "test_acc = runner.test()\n",
    "\n",
    "runtime = time.time() - t0\n",
    "print('Total runtime: {:.2f}s'.format(runtime))\n",
    "runner.logger.write('Total runtime: {:.2f}s\\n'.format(runtime))\n",
    "runner.logger.close()\n",
    "\n",
    "print('val acc: {}, test acc: {}'.format(val_acc, test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c537b575",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
